package com.wxweven.algorithm.array.presum;

import java.util.HashMap;
import java.util.Map;

class LC560和为k的子数组个数 {

    //暴力解法，3重循环，时间复杂度 O(n3)
    public int solution1(int[] nums, int k) {
        int count = 0;
        for (int i = 0; i < nums.length; i++) {
            for (int j = i; j < nums.length; j++) {
                int sum = 0;
                for (int index = i; index <= j; index++) {
                    sum += nums[index];
                }

                if (sum == k) {
                    count++;
                }
            }
        }

        return count;
    }

    //2重循环，利用前缀和，时间复杂度 O(n2)
    public int solution2(int[] nums, int k) {
        int count = 0;
        for (int i = 0; i < nums.length; i++) {
            int prefixSum = 0;
            for (int j = i; j < nums.length; j++) {
                prefixSum += nums[j];

                if (prefixSum == k) {
                    count++;
                }
            }
        }

        return count;
    }

    /**
     * 视频讲解：https://www.youtube.com/watch?v=aYfwus5T3Bs
     */
    public int solution3(int[] nums, int k) {
        if (nums == null || nums.length == 0) {
            return 0;
        }

        // 更高效的解法
        int count = 0;
        int prefixSum = 0;

        Map<Integer, Integer> map = new HashMap<>();
        // 这个是初始条件，前缀和数组会比正常的数组多一个和为0的
        map.put(0, 1);

        /*
         * 假设 nums = [1, 3, 5, 4, -1, 1]
         * 则 sum  = [0, 1, 4, 9, 13, 12, 13]
         *
         * sum[i...j] = prefixSum[j] - prefixSum[i-1]
         * 证明：
         * prefixSum[j] = nums[0]+nums[1]+...+nums[i-1]+nums[i]+...+ nums[j]
         * prefixSum[i-1] = nums[0]+nums[1]+...+nums[i-1]
         * prefixSum[j] - prefixSum[i-1] = nums[i]+...+ nums[j] = sum[i...j]
         *
         * 问题转化为，在sum数组中，有多种<i, j> (j>=i)下标对，使得 sum[j] - sum[i] = k, 即 sum[i] = sum[j] - k
         *
         * Map<Integer, Integer> map，map中的key是前缀和，value是前缀和次数
         * 于是可以遍历sum数组，遍历到sum[i]时，找下map中有没有key为 sum[j] - k 的，如果有，则说明又找到了一种解
         * 举个例子，sum  = [0, 1, 4, 9, 13, 12, 13]，假设k=4，
         *  当前遍历到 1，那就是看map中key为-3的有几个，发现没有，count不变；
         *  当前遍历到 4，那就是看map中key为0的有几个，发现只有1个，count+1；
         *  当前遍历到 9，那就是看map中key为5的有几个, 发现没有，count不变；
         *  当前遍历到 13，那就是看map中key为13-4=9的有几个, 发现有1个，count+1；
         *  当前遍历到 12，那就是看map中key为12-4=8的有几个, 发现没有，count不变；
         *  当前遍历到 13，那就是看map中key为13-4=9的有几个, 发现有1个，count+1；
         *
         *  最终，count=3
         */
        for (int num : nums) {
            prefixSum += num;

            count += map.getOrDefault(prefixSum - k, 0);

            map.put(prefixSum, map.getOrDefault(prefixSum, 0) + 1);
        }

        return count;
    }
}